package com.kvn.dubbo.filter;

import com.alibaba.dubbo.common.Constants;
import com.alibaba.dubbo.common.extension.Activate;
import com.alibaba.dubbo.rpc.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.UUID;

/**
 * Created by Administrator on 2020/3/2 0002.
 */
@Activate(group = {Constants.PROVIDER, Constants.CONSUMER})
public class DemoFilter implements Filter {
    private static final Logger log = LoggerFactory.getLogger(DemoFilter.class);
    private static final String TRACE_ID = "traceId";


    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        boolean providerSide = RpcContext.getContext().isProviderSide();

        String traceId = MDC.get(TRACE_ID);
        // 1. 如果当前线程中有没有 traceId，则生成一个 traceId，放到当前线程中
        if (traceId == null) {
            // 服务提供者，首先从消费端获取 traceId
            if (providerSide) {
                traceId = RpcContext.getContext() != null ? RpcContext.getContext().getAttachment(TRACE_ID) : null;
            }

            if (traceId == null) {
                traceId = UUID.randomUUID().toString();
            }
            MDC.put(TRACE_ID, traceId);
        }

        // 2. 如果是服务消费者，则将 traceId 往服务提供者传递
        if (RpcContext.getContext().isConsumerSide()) {
            RpcContext.getContext().setAttachment(TRACE_ID, traceId);
        }

        log.info("--->traceId:{}", traceId);
        try {
            return invoker.invoke(invocation);
        } finally {
            // 3. 如果是服务提供者，调用完成后，将 uid 清除
            if (providerSide) {
                MDC.remove(TRACE_ID);
            }
        }
    }
}